from collections.abc import Iterable, Mapping, Sequence
from typing import TypeVar

from zkl_aiutils_datasets.basics import Dataset
from zkl_aiutils_datasets.processing import ChainedDatasets, DelegatingDataset, KnownSizeChainedDatasets
from .simple import SimpleGroupedDatasets, SimpleIndexedDatasets, SimpleIndexedNamedDatasets, SimpleNamedDatasets

AnySample = TypeVar('AnySample')
AnyDataset = TypeVar('AnyDataset', bound=Dataset)


class ChainedGroupedDatasets(DelegatingDataset[AnySample], SimpleGroupedDatasets[AnyDataset]):
    def __init__(self, children: Iterable[AnyDataset]):
        SimpleGroupedDatasets.__init__(self, children)
        DelegatingDataset.__init__(self, ChainedDatasets(children))


class ChainedIndexedDatasets(DelegatingDataset[AnySample], SimpleIndexedDatasets[AnyDataset]):
    def __init__(self, indexed_children: Sequence[AnyDataset]):
        SimpleIndexedDatasets.__init__(self, indexed_children)
        DelegatingDataset.__init__(self, KnownSizeChainedDatasets(indexed_children))


class ChainedNamedDatasets(DelegatingDataset[AnySample], SimpleNamedDatasets[AnyDataset]):
    def __init__(self, named_children: Mapping[str, AnyDataset]):
        SimpleNamedDatasets.__init__(self, named_children)
        DelegatingDataset.__init__(self, KnownSizeChainedDatasets(named_children.values()))


class ChainedIndexedNamedDatasets(DelegatingDataset[AnySample], SimpleIndexedNamedDatasets[AnyDataset]):
    def __init__(self, children: Mapping[str, Dataset[AnySample]] | Iterable[tuple[str, Dataset[AnySample]]]):
        SimpleIndexedNamedDatasets.__init__(self, children)
        DelegatingDataset.__init__(self, KnownSizeChainedDatasets(self.indexed_children))
